import { useState, useRef, useCallback } from 'react';

export interface ChatMessage {
  id: string;
  role: 'user' | 'assistant';
  content: string;
  contents?: {
    type: 'text' | 'thinking' | 'command' | 'output' | 'system';
    text: string;
    language?: string;
    sayType?: string;
  }[];
  timestamp: number;
  isStreaming?: boolean;
}

export const useChatStream = (unitId: string) => {
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [isStreaming, setIsStreaming] = useState(false);
  const eventSourceRef = useRef<EventSource | null>(null);
  const currentMessageRef = useRef<ChatMessage | null>(null);

  const sendMessage = useCallback((userMessage: string) => {
    const user: ChatMessage = { id: `user-${Date.now()}`, role: 'user', content: userMessage, timestamp: Date.now() };
    setMessages(prev => [...prev, user]);

    const assistant: ChatMessage = { id: `assistant-${Date.now()}`, role: 'assistant', content: '', contents: [], timestamp: Date.now(), isStreaming: true };
    currentMessageRef.current = assistant;
    setMessages(prev => [...prev, assistant]);
    setIsStreaming(true);

    const url = `/api/agent/stream?unitId=${encodeURIComponent(unitId)}&message=${encodeURIComponent(userMessage)}`;
    const es = new EventSource(url);
    eventSourceRef.current = es;

    es.addEventListener('message', (event) => {
      try {
        const data = JSON.parse((event as MessageEvent).data);
        const text: string = data.text || '';
        const sayType: string = data.say || data.type || 'text';

        let contentType: 'text' | 'thinking' | 'system';
        switch (sayType) {
          case 'reasoning':
            contentType = 'thinking';
            break;
          case 'completion_result':
            // 作为最终答案的主体内容展示
            contentType = 'text';
            break;
          case 'task':
          case 'task_progress':
          case 'api_req_started':
          case 'api_req_finished':
          case 'api_req_retried':
          case 'api_req_retry':
          case 'user_feedback':
          case 'user_feedback_diff':
          case 'api_req_retry_scheduled':
          case 'mcp_server_request_started':
          case 'mcp_server_response':
          case 'mcp_notification':
          case 'use_mcp_server':
          case 'load_mcp_documentation':
          case 'tool':
          case 'browser_action_launch':
          case 'browser_action':
          case 'browser_action_result':
          case 'shell_integration_warning':
          case 'shell_integration_warning_with_suggestion':
          case 'deleted_api_reqs':
          case 'diff_error':
          case 'clineignore_error':
          case 'info':
          case 'error':
          case 'error_retry':
          case 'checkpoint_created':
            contentType = 'system';
            break;
          default:
            contentType = 'text';
        }

        if (currentMessageRef.current) {
          currentMessageRef.current.contents!.push({
            type: contentType,
            text,
            sayType,
          });
          if (contentType === 'text' || contentType === 'thinking') {
            currentMessageRef.current.content += text;
          }
          setMessages(prev => [...prev]);
        }
      } catch (e) {
        // ignore
      }
    });

    es.addEventListener('command', (event) => {
      try {
        const data = JSON.parse((event as MessageEvent).data);
        const text = data.text || '';
        if (currentMessageRef.current) {
          currentMessageRef.current.contents!.push({
            type: 'command',
            text,
            language: 'bash',
            sayType: data.say || 'command',
          });
          setMessages(prev => [...prev]);
        }
      } catch {}
    });

    es.addEventListener('command_output', (event) => {
      try {
        const data = JSON.parse((event as MessageEvent).data);
        const text = data.text || '';
        if (currentMessageRef.current) {
          currentMessageRef.current.contents!.push({
            type: 'output',
            text,
            sayType: data.say || 'command_output',
          });
          setMessages(prev => [...prev]);
        }
      } catch {}
    });

    es.addEventListener('ask', (event) => {
      try {
        const data = JSON.parse((event as MessageEvent).data);
        const text = data.text || '';
        if (currentMessageRef.current) {
          currentMessageRef.current.contents!.push({
            type: 'output',
            text,
            sayType: data.say || 'ask',
          });
          setMessages(prev => [...prev]);
        }
      } catch {}
    });

    es.addEventListener('complete', () => {
      es.close();
      setIsStreaming(false);
      if (currentMessageRef.current) {
        currentMessageRef.current.isStreaming = false;
        setMessages(prev => [...prev]);
      }
    });

    es.addEventListener('error', () => {
      es.close();
      setIsStreaming(false);
      if (currentMessageRef.current) {
        currentMessageRef.current.isStreaming = false;
        setMessages(prev => [...prev]);
      }
    });
  }, [unitId]);

  const stopStreaming = useCallback(() => {
    if (eventSourceRef.current) {
      eventSourceRef.current.close();
      eventSourceRef.current = null;
    }
    setIsStreaming(false);
    if (currentMessageRef.current) {
      currentMessageRef.current.isStreaming = false;
      setMessages(prev => [...prev]);
    }
  }, []);

  return { messages, isStreaming, sendMessage, stopStreaming };
};
